package com.beamfield.common;

import com.beamfield.common.CommonUtils;
import com.beamfield.common.ExcelUtil;
import com.beamfield.common.IKAnalyzer4Lucene7;
import com.beamfield.modelVo.DocumentVo;
import org.apache.commons.lang3.StringUtils;
import org.apache.lucene.analysis.TokenStream;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.StringField;
import org.apache.lucene.document.TextField;
import org.apache.lucene.index.DirectoryReader;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.index.Term;
import org.apache.lucene.queryparser.classic.QueryParser;
import org.apache.lucene.search.*;
import org.apache.lucene.search.highlight.*;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FSDirectory;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.text.PDFTextStripper;
import org.apache.poi.POITextExtractor;
import org.apache.poi.hslf.extractor.PowerPointExtractor;
import org.apache.poi.hssf.usermodel.HSSFWorkbook;
import org.apache.poi.hwpf.extractor.WordExtractor;
import org.apache.poi.ss.usermodel.Row;
import org.apache.poi.ss.usermodel.Sheet;
import org.apache.poi.ss.usermodel.Workbook;
import org.apache.poi.xssf.usermodel.XSSFWorkbook;
import org.apache.poi.xwpf.extractor.XWPFWordExtractor;
import org.apache.poi.xwpf.usermodel.XWPFDocument;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.core.env.Environment;
import org.springframework.stereotype.Component;

import java.io.*;
import java.nio.file.Paths;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Properties;
import java.util.stream.Collectors;

/**
 * @author Zephyr
 * @Description:
 * @Date 2018/9/25
 */
@Component
public class LuceneUtils {

    public static Directory dir;

    public static Properties properties = new Properties();

    static {
        try {
            String path = "config.properties";
            InputStream inStream = CommonUtils.class.getClassLoader().getResourceAsStream(path);
            if (inStream == null) {
                inStream = CommonUtils.class.getClassLoader().getResourceAsStream("/" + path);
            }
            properties.load(inStream);

            dir = FSDirectory.open(Paths.get(properties.getProperty("luceneDir")));
        } catch (IOException e) {
            e.printStackTrace();
        }
    }

    /**
     * pdf文件转string
     * @param file
     * @return
     */
    public static String readPdf(File file){
        StringBuilder content = new StringBuilder();
        PDDocument pdDocument = null;
        try {
            pdDocument = PDDocument.load(file);
            PDFTextStripper ts = new PDFTextStripper();
            content.append(ts.getText(pdDocument));
            pdDocument.close();
            return content.toString().trim();
        }catch (Exception e){
            e.printStackTrace();
        }finally {
            try {
                pdDocument.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return "";
    }

    /**
     * txt格式文件转string
     * @param file
     * @return
     */
    public static String readTxt(File file) {
        StringBuilder content = new StringBuilder();
        try {
            InputStreamReader reader = new InputStreamReader(new FileInputStream(file), "UTF-8");
            BufferedReader br = new BufferedReader(reader);
            String s1 = null;

            while ((s1 = br.readLine()) != null) {
                content.append(s1).append("\r");
            }
            br.close();
            reader.close();
        } catch (IOException e) {
            e.printStackTrace();
        }
        return content.toString().trim();
    }


    /**
     * 获取IndexWriter实例
     * @return
     */
    private IndexWriter getWriter() {
        try {
            IndexWriterConfig iwc = new IndexWriterConfig(new IKAnalyzer4Lucene7(true));
            return new IndexWriter(dir, iwc);
        }catch (Exception e){
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 添加索引
     * @param id 文件数据库id
     * @param fileNmae 文件名称
     * @param fileType 文件类型
     * @param file 文件
     */
    public void addIndex(String id, String fileNmae, String fileType, File file, Long beamId)  {
        IndexWriter indexWriter = this.getWriter();
        if(indexWriter == null){
            return ;
        }
        try {
            String content = getContent(file);
            Document document = new Document();
            document.add(new StringField("id", id, Field.Store.YES));
            document.add(new StringField("fileType", fileType, Field.Store.YES));
            document.add(new StringField("beamId", String.valueOf(beamId), Field.Store.YES));
            document.add(new TextField("fileName", fileNmae, Field.Store.YES));
            document.add(new TextField("content", content, Field.Store.YES));
            indexWriter.addDocument(document);
        }catch (Exception e){
            e.printStackTrace();
        }finally {
            try {
                indexWriter.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }

    /**
     *  根据文件类型获取文件内容
     * @param file
     * @return
     */
    private String getContent(File file) {
        String fileName = file.getName();
        String type = fileName.substring(fileName.lastIndexOf(".")+1, fileName.length());
        switch (type){
            case "pdf": return readPdf(file);
            case "txt": return readTxt(file);
            case "doc": return readDoc(file);
            case "docx": return readDocx(file);
            case "xls": return readXls(file);
            case "xlsx": return readXlsx(file);
            case "ppt": return readPpt(file);
            default: return "";
        }
    }

    /**
     * ppt格式转string
     * @param file
     * @return
     */
    private String readPpt(File file) {
        StringBuilder sb = new StringBuilder();
        PowerPointExtractor extractor = null;
        try {
            extractor = new PowerPointExtractor(new FileInputStream(file));
            sb.append(extractor.getText());
        } catch (IOException e) {
            e.printStackTrace();
        }finally {
            try {
                if(extractor != null) {
                    extractor.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return sb.toString();
    }

    /**
     * xlsx格式的excel转string
     * @param file
     * @return
     */
    private String readXlsx(File file) {
        StringBuilder sb = new StringBuilder();
        Workbook workbook = null;
        try {
            workbook = new XSSFWorkbook(new FileInputStream(file));
            for(int i = 0, j = workbook.getNumberOfSheets(); i < j; i++){
                Sheet sheet = workbook.getSheetAt(i);
                for(int k = 0, l = sheet.getLastRowNum(); k < l; k++){
                    Row row = sheet.getRow(k);
                    for(int m = 0, n = row.getLastCellNum(); m < n; m++){
                        String value = ExcelUtil.getCellValueAsString(row.getCell(m));
                        if(StringUtils.isNotEmpty(value)){
                            sb.append(value);
                        }
                    }
                }
            }
        }catch (Exception e){
            e.printStackTrace();
        }finally {
            try {
                if(workbook != null) {
                    workbook.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return sb.toString();
    }

    /**
     * xls 格式的excel转string
     * @param file
     * @return
     */
    private String readXls(File file) {
        StringBuilder sb = new StringBuilder();
        Workbook workbook = null;
        try {
            workbook = new HSSFWorkbook(new FileInputStream(file));
            for(int i = 0, j = workbook.getNumberOfSheets(); i < j; i++){
                Sheet sheet = workbook.getSheetAt(i);
                for(int k = 0, l = sheet.getLastRowNum(); k < l; k++){
                    Row row = sheet.getRow(k);
                    for(int m = 0, n = row.getLastCellNum(); m < n; m++){
                        String value = ExcelUtil.getCellValueAsString(row.getCell(m));
                        if(StringUtils.isNotEmpty(value)){
                            sb.append(value);
                        }
                    }
                }
            }
        }catch (Exception e){
            e.printStackTrace();
        }finally {
            try {
                if(workbook != null) {
                    workbook.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return sb.toString();
    }

    /**
     *  docx格式的word转string
     * @param file
     * @return
     */
    private String readDocx(File file) {
        StringBuilder sb = new StringBuilder();
        POITextExtractor extractor = null;
        try {
            XWPFDocument xwpfDocument = new XWPFDocument(new FileInputStream(file));
            extractor = new XWPFWordExtractor(xwpfDocument);
            sb.append(extractor.getText());
        } catch (IOException e) {
            e.printStackTrace();
        }finally {
            try {
                if(extractor != null) {
                    extractor.close();
                }
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
        return sb.toString();
    }

    /**
     * word 转string
     * @param file
     * @return
     */
    private String readDoc(File file) {
        StringBuilder sb = new StringBuilder();
        WordExtractor extractor = null;
        try {
            extractor = new WordExtractor(new FileInputStream(file));
            sb.append(extractor.getText());
        }catch (Exception e){
            e.printStackTrace();
        }finally {
            try {
                if(extractor != null){
                    extractor.close();
                }
            }catch (IOException e){
                e.printStackTrace();
            }
        }
        return sb.toString();
    }

    /**
     * 查询索引
     * @param fileType
     * @param queryWords
     * @return
     */
    public List<DocumentVo> search(String fileType, String queryWords, Long beamId) {
        DirectoryReader directoryReader = null;
        try {
            directoryReader = DirectoryReader.open(dir);
            IndexSearcher indexSearcher = new IndexSearcher(directoryReader);

            BooleanQuery.Builder booleanQueryBuilder = new BooleanQuery.Builder();

            IKAnalyzer4Lucene7 analyzer = new IKAnalyzer4Lucene7(true);


            BooleanQuery booleanQuery = booleanQueryBuilder.add(new WildcardQuery(new Term("fileName", "*" + queryWords + "*")), BooleanClause.Occur.SHOULD)
                    .add(new QueryParser("content", analyzer).parse(queryWords), BooleanClause.Occur.SHOULD)
                    .add(new WildcardQuery(new Term("content", "*" + queryWords + "*")), BooleanClause.Occur.SHOULD).build();


            booleanQueryBuilder = new BooleanQuery.Builder();

            boolean flag = false;

            if(StringUtils.isNotEmpty(fileType)){

                booleanQueryBuilder.add(new TermQuery(new Term("fileType", fileType)), BooleanClause.Occur.MUST);
                flag = true;
            }

            if(beamId != null && beamId != 0L){

                booleanQueryBuilder.add(new TermQuery(new Term("beamId", String.valueOf(beamId))), BooleanClause.Occur.MUST);
                flag = true;
            }

            if(flag){
                booleanQueryBuilder.add(booleanQuery, BooleanClause.Occur.MUST);

                booleanQuery = booleanQueryBuilder.build();
            }

            TopDocs hits = indexSearcher.search(booleanQuery, 100);
            SimpleHTMLFormatter simpleHTMLFormatter = new SimpleHTMLFormatter("<span style='color:red'>", "</span>");
            Highlighter highlighter = new Highlighter(simpleHTMLFormatter, new QueryScorer(booleanQuery));
            highlighter.setTextFragmenter(new SimpleFragmenter(30));
            return Arrays.stream(hits.scoreDocs).map(scoreDoc -> {
                try {
                    DocumentVo documentVo = new DocumentVo();
                    Document doc = indexSearcher.doc(scoreDoc.doc);
                    String content = doc.get("content");
                    String fileName = doc.get("fileName");
                    // 内容增加高亮显示
                    content = highlighter.getBestFragment(analyzer, "content", content);
                    if(StringUtils.isNotEmpty(content)){
                        documentVo.setHighLightContent(content);
                    }
                    fileName = highlighter.getBestFragment(analyzer, "fileName", fileName);
                    if(StringUtils.isNotEmpty(fileName)){
                        documentVo.setFormer_name(fileName);
                    }
                    documentVo.setClassifyId(Long.parseLong(doc.get("id")));
                    return documentVo;
                } catch (IOException | InvalidTokenOffsetsException e) {
                    e.printStackTrace();
                }                return null;
            }).collect(Collectors.toList());

        } catch (Exception e) {
            e.printStackTrace();
        }finally {
            try {
                if(directoryReader != null){
                    directoryReader.close();
                }
            }catch (IOException e){
                e.printStackTrace();
            }

        }
        return new ArrayList<>();
    }

    /**
     * 删除索引
     * @param id
     */
    public void delIndex(String id){
        IndexWriter indexWriter = this.getWriter();
        if(indexWriter == null){
            return ;
        }
        try {
            indexWriter.deleteDocuments(new Term("id", id));
        } catch (IOException e) {
            e.printStackTrace();
        }finally {
            try {
                indexWriter.close();
            } catch (IOException e) {
                e.printStackTrace();
            }
        }
    }
}
